Add configurable MoE runtime flags to MegatronConfig#1213
Draft
tyler-griggs wants to merge 2 commits intomainfrom
Draft
Add configurable MoE runtime flags to MegatronConfig#1213tyler-griggs wants to merge 2 commits intomainfrom
tyler-griggs wants to merge 2 commits intomainfrom
Conversation
Expose 5 MoE runtime config flags as first-class MegatronConfig fields: - moe_token_dispatcher_type (replaces hardcoded "alltoall") - moe_router_load_balancing_type (replaces hardcoded "none") - moe_grouped_gemm (enables fused grouped GEMM for MoE) - moe_router_score_function (e.g. "sigmoid" for GLM/DeepSeek-V3) - moe_router_enable_expert_bias (learned bias for load balancing) Most architecture flags (num_experts, topk, qkv_bias, rotary, etc.) are auto-detected by AutoBridge from the HF config.json. These 5 flags control runtime behavior that is NOT in the HF config. Advanced/model-specific flags can still be passed through the existing transformer_config_kwargs dict, which is applied after these fields. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This was referenced Feb 25, 2026
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request correctly exposes several Mixture-of-Experts (MoE) runtime flags in MegatronConfig, replacing previous hardcoded values. This improves configurability. The changes are well-implemented across the configuration files, worker logic, and are accompanied by new tests.
My review includes a couple of suggestions for improvement:
- In
megatron_worker.py, I've suggested a small refactoring to make the application of optional MoE configurations more concise and maintainable. - In the new test file
test_moe_config.py, I've pointed out a test case that could be made more exhaustive by adding assertions for all the fields being tested.
Overall, this is a solid contribution that enhances the flexibility of MoE configurations.
Comment on lines
+244
to
+247
| if megatron_config.moe_router_score_function is not None: | ||
| provider.moe_router_score_function = megatron_config.moe_router_score_function | ||
| if megatron_config.moe_router_enable_expert_bias is not None: | ||
| provider.moe_router_enable_expert_bias = megatron_config.moe_router_enable_expert_bias |
Contributor
There was a problem hiding this comment.
To improve maintainability and reduce repetition, you can loop over the optional MoE configuration fields. This makes it easier to add more optional fields in the future.
Suggested change
| if megatron_config.moe_router_score_function is not None: | |
| provider.moe_router_score_function = megatron_config.moe_router_score_function | |
| if megatron_config.moe_router_enable_expert_bias is not None: | |
| provider.moe_router_enable_expert_bias = megatron_config.moe_router_enable_expert_bias | |
| for field_name in ("moe_router_score_function", "moe_router_enable_expert_bias"): | |
| value = getattr(megatron_config, field_name) | |
| if value is not None: | |
| setattr(provider, field_name, value) |
Comment on lines
+49
to
+51
| assert cfg.moe_grouped_gemm is True | ||
| assert cfg.moe_router_score_function == "sigmoid" | ||
| assert cfg.moe_router_enable_expert_bias is True |
Contributor
There was a problem hiding this comment.
This test is not exhaustive. It's good practice to assert all fields that are being set from the dictionary to ensure the build_nested_dataclass function works as expected for all new fields. The assertions for moe_token_dispatcher_type and moe_router_load_balancing_type are missing.
Suggested change
| assert cfg.moe_grouped_gemm is True | |
| assert cfg.moe_router_score_function == "sigmoid" | |
| assert cfg.moe_router_enable_expert_bias is True | |
| assert cfg.moe_token_dispatcher_type == "alltoall" | |
| assert cfg.moe_router_load_balancing_type == "none" | |
| assert cfg.moe_grouped_gemm is True | |
| assert cfg.moe_router_score_function == "sigmoid" | |
| assert cfg.moe_router_enable_expert_bias is True |
…oad_balancing_type in test_moe_config_from_dict Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Exposes 5 MoE runtime config flags as first-class fields on the
MegatronConfigdataclass, replacing 2 hardcoded values in the Megatron worker.New fields
moe_token_dispatcher_typestr"alltoall"moe_router_load_balancing_typestr"none"moe_grouped_gemmboolFalsemoe_router_score_functionOptional[str]None"sigmoid"for GLM/DeepSeek-V3)moe_router_enable_expert_biasOptional[bool]Noneconfig.json— no explicit config needed